import pymongo
import sys
from optparse import OptionParser

#	This is pyVOMS
#
#       v4: Argument and option handling cleaned up!
#       v3: This also manages a site list of cloud sites participating in each VO


def connect_mongo(dbname,collname):
    try:
        mongo_conn = pymongo.Connection('ec2-54-234-85-121.compute-1.amazonaws.com',27017)
        db = mongo_conn['mvom']
        coll=db[collname]
        return coll;
    except:
        print 'Unable to connect to the pyVOMS server.'
        sys.exit(-1)


def stat_voms( opts ):
    if  opts.vo_name == None:
        #  Stat the VOMS and list just the existing VO names
        collection = connect_mongo("pvom","vom_list")
        results = collection.find({'vo_name':{'$regex':'^'}})
        print "Current VOs:"
        print "============" 
        vo_list = []
        for r in results:
            if r['vo_name'] not in vo_list:
                vo_list.append( r['vo_name'] )
                print r['vo_name']
    else:
        #  Stat the VOMS for a specific VO and list its sites and members
        collection = connect_mongo("pvom","site_list")
        results = collection.find({'vo_name':opts.vo_name})
        if  results.count() == 0:
            print "VO not found or malformed.  There must be at least one participating site in a VO.\n"
            sys.exit(-1)
        
        print opts.vo_name, ": Participating Sites"
        for r in results:
            print "  Cloud Name:", r['cloud_name'], " AUTH_URL:", r['auth_url']

        collection = connect_mongo("pvom","vom_list")
        results = collection.find({'vo_name':opts.vo_name})
        print opts.vo_name, ": Members"
        for r in results:
            #print r
            if  len(r) == 8:
                print "  User Name:", r['user'], " User Pswd:", r['passwd'], " VO Role:", r['role'], " Role Pswd:", r['role_pswd']
            if  len(r) == 7:
                print "  User Name:", r['user'], " User Pswd:", r['passwd'], " VO Role:", r['role']
            if  len(r) == 9 or len(r) == 6:
                print "  Cloud Name:", r['cloud_name'], " User Name:", r['user'], " VO Role:", r['role'], " Role Pswd:", r['role_pswd']


def add_user( opts ):
    ## add a user to a VO 
    if opts.user == None or opts.cloud_name == None or opts.vo_name == None or opts.role == None or opts.role_pswd == None:
        print "Usage: pyVOMS adduser -v VOname -c CloudName -u username -r role -p role_password\n"
        sys.exit(-1) 
    vo_memb_dict = { 'vo_name':opts.vo_name, 'cloud_name':opts.cloud_name, 'user':opts.user, 'role':opts.role, 'role_pswd':opts.role_pswd }
    try:
        collection = connect_mongo("pvom","vom_list")
        collection.save( vo_memb_dict )
    except:
        print "Error adding user.\n"
        sys.exit(-1)
    print "  User <", opts.user, "> from Cloud <", opts.cloud_name, "> added to VO <", opts.vo_name, ">"


def add_site( opts ):
    ### Add a cloud site to a VO
    if opts.cloud_name == None or opts.vo_name == None or opts.auth == None:
        print "Usage: pyVOMS addsite -v VOname -c CloudName -a AUTH_URL\n"
        sys.exit(-1)
    site_dict = { 'cloud_name':opts.cloud_name, 'vo_name':opts.vo_name, 'auth_url':opts.auth }
    try:
        collection = connect_mongo("pvom","site_list")
        collection.save( site_dict )
    except:
        print "Error adding site.\n"
        sys.exit(-1)
    print "Site ", opts.cloud_name, " added to VO", opts.vo_name


def del_user( opts ):
    ### delete a user from a VO
    if opts.vo_name == None or opts.cloud_name == None or opts.user == None:
        print "Usage: pyVOMS deluser -v VOname -c CloudName -u username\n"
        sys.exit(-1)
    vo_memb_dict = { 'vo_name':opts.vo_name, 'cloud_name':opts.cloud_name, 'user':opts.user }
    try: 
        collection = connect_mongo("pvom","vom_list")
        collection.remove( vo_memb_dict )
    except:
        print "Error deleting user.\n"
        sysexit(-1)
    print "  User <", opts.user, "> from Cloud <", opts.cloud_name, "> deleted from VO <", opts.vo_name, ">"


def del_site( opts ):
    ### Delete a cloud site from a VO
    if opts.vo_name == None or opts.cloud_name == None:
        print "Usage:  pyVOMS delsite -v VOname -c CloudName\n"
        sys.exit(-1)
    site_dict = { 'vo_name':opts.vo_name, 'cloud_name':opts.cloud_name }

    # Delete all users from this cloud first
    try: 
        collection = connect_mongo( "pvom", "vom_list" )
        collection.remove( site_dict )
    except:
        print 'Error deleting users from a site.\n'
        sys.exit(-1)

    # Now delete site from site_list
    try: 
        collection = connect_mongo("pvom","site_list")
        collection.remove( site_dict )
    except:
        print 'Error deleting site.\n'
        sys.exit(-1)
    print "Site (and users) deleted."


def del_vo( opts ):
    ### Delete an entire VO
    if opts.vo_name == None:
        print "Usage: pyVOMS delvo -v vo_name\n"
        sys.exit(-1)
    vo_dict = { 'vo_name':opts.vo_name }

    # Delete all VO users from vom_list
    try: 
        collection = connect_mongo( "pvom", "vom_list" )
        collection.remove( vo_dict )
    except:
        print "Error deleting users from a VO.\n"
        sys.exit(-1)

    # Now delete VO from site_list
    try: 
        collection = connect_mongo("pvom","site_list")
        collection.remove( vo_dict )
    except:
        print "Error deleting VO.\n"
        sys.exit(-1)
    print "VO deleted."


if  __name__ == '__main__':

#    Usage
#
#    pyVOMS [command] [options]
#
#    stat
#    stat -v VOname
#    adduser -v VOname -c CloudName -u username (? -k user_password) -r role -p role_password
#    addsite -v VOname -c CloudName -a AUTH_URL
#    deluser -v VOname -c CloudName -u username
#    delsite -v VOname -c CloudName
#    delvo -v VOname

    ##  Parse the command line 

    usage = "Usage: %prog command [options]\n\
    stat\n\
    stat -v VOname\n\
    adduser -v VOname -c CloudName -u username -r role -p role_password\n\
    addsite -v VOname -c CloudName -a AUTH_URL\n\
    deluser -v VOname -c CloudName -u username\n\
    delsite -v VOname -c CloudName\n\
    delvo -v VOname"

    parser = OptionParser(usage=usage)
    parser.add_option("-v", dest="vo_name", help="VO name for this operation")
    parser.add_option("-c", dest="cloud_name", help="Name of cloud participating in this VO")
    parser.add_option("-u", dest="user", help="username")
    parser.add_option("-k", dest="passwd", help="user password")
    parser.add_option("-r", dest="role", help="Role for this VO member (username)")
    parser.add_option("-p", dest="role_pswd", help="Role password for this VO member")
    parser.add_option("-a", dest="auth", help="AUTH_URL for adding a site to a VO") 

    (options, args) = parser.parse_args()
    if  len(args) != 1 :
        print "Exactly one command arg, please, or use --help if you need help.\n"
        sys.exit(-1)
    command = args[0]

#   Execute the command

    if  command == 'stat':
        stat_voms( options )
    elif  command == 'adduser':
        add_user( options )
    elif  command == 'addsite':
        add_site( options )
    elif  command == 'deluser':
        del_user( options )
    elif  command == 'delsite':
        del_site( options )
    elif  command == 'delvo':
        del_vo( options )
    else:
        print "Unrecognized command.\n"
        sys.exit(-1)

    sys.exit(0)
